#!/usr/bin/env python3
"""
A股每日数据更新脚本
用于更新所有A股股票的最新数据
"""

import sys
import os
import time
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Dict, Any
import json
import pandas as pd

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.data_source_manager import DataSourceManager
from src.data_storage import DataStorage
from src.utils.config import get_settings, ensure_data_dirs
from src.utils.logger import get_logger

logger = get_logger("AShareUpdate")


class AShareUpdater:
    """A股数据更新器"""

    def __init__(self):
        self.settings = get_settings()
        self.logger = logger
        self.data_source_manager = DataSourceManager(auto_check=False)
        self.data_storage = DataStorage()
        self.update_config = self._load_update_config()

    def _load_update_config(self) -> Dict[str, Any]:
        """加载更新配置"""
        return {
            'batch_size': 100,  # 每批处理的股票数量
            'delay_between_batches': 1.0,  # 批次间延迟（秒）
            'delay_between_stocks': 0.05,  # 股票间延迟（秒）
            'max_retries': 3,  # 最大重试次数
            'days_to_update': 7,  # 更新最近N天的数据
            'interval': '1d',  # 数据间隔
            'skip_no_data': True,  # 跳过无数据的股票
            'filter_inactive': True,  # 过滤非活跃股票
            'save_report': True,  # 保存更新报告
            'parallel_requests': 5,  # 并发请求数量
            'update_failed_stocks': True,  # 重试之前失败的股票
        }

    def get_update_stock_list(self, update_mode: str = "incremental") -> List[str]:
        """获取需要更新的股票列表"""
        try:
            self.logger.info("获取股票列表...")

            # 使用数据源管理器获取合并的股票列表
            df = self.data_source_manager.get_combined_stock_list()

            if df.empty:
                self.logger.warning("未获取到任何股票列表")
                return []

            self.logger.info(f"获取到 {len(df)} 只股票（来自多个数据源）")

            # 过滤非活跃股票
            if 'status' in df.columns and self.update_config['filter_inactive']:
                df = df[df['status'] == '0']  # status='0'表示上市交易
                self.logger.info(f"过滤非活跃股票后剩余: {len(df)} 只")

            # 过滤新三板（通常流动性较差）
            if 'market_type' in df.columns:
                df = df[df['market_type'] != '新三板']

            # 过滤可转债和其他非股票产品
            if 'type' in df.columns:
                stock_df = df[df['type'] == '1']  # type=1通常是股票
                filtered_count = len(df) - len(stock_df)
                if filtered_count > 0:
                    self.logger.info(f"过滤可转债和其他非股票产品: {filtered_count} 只")
                df = stock_df

            # 额外过滤：只保留主板、创业板、科创板的股票
            if 'market_type' in df.columns:
                valid_markets = ['深圳主板', '上海主板', '创业板', '科创板']
                market_df = df[df['market_type'].isin(valid_markets)]
                if len(market_df) < len(df):
                    filtered_count = len(df) - len(market_df)
                    self.logger.info(f"过滤其他市场类型: {filtered_count} 只")
                df = market_df

            if update_mode == "full":
                # 全量更新：返回所有符合条件的股票
                symbols = df['symbol'].tolist()
                self.logger.info(f"全量更新：获取 {len(symbols)} 只股票")

                # 按数据源分组统计
                if 'data_source' in df.columns:
                    source_groups = df.groupby('data_source')
                    for source, group in source_groups:
                        self.logger.info(f"  {source}: {len(group)} 只")

                # 按市场类型分组
                if 'market_type' in df.columns:
                    market_groups = df.groupby('market_type')
                    for market, group in market_groups:
                        self.logger.info(f"  {market}: {len(group)} 只")

                return symbols
            else:
                # 增量更新：只返回有历史数据的股票
                stocks_with_data = []
                for symbol in df['symbol'].tolist():
                    info = self.data_storage.get_data_info(symbol, self.update_config['interval'])
                    if info['file_exists'] and info['record_count'] > 0:
                        stocks_with_data.append(symbol)

                self.logger.info(f"增量更新：找到 {len(stocks_with_data)} 只有历史数据的股票需要更新")

                return stocks_with_data

        except Exception as e:
            self.logger.error(f"获取更新股票列表失败: {e}")
            return []

    def update_stock(self, symbol: str, retry_count: int = 0, update_mode: str = "incremental") -> Dict[str, Any]:
        """更新单个股票数据"""
        result = {
            'symbol': symbol,
            'success': False,
            'new_records': 0,
            'error': None,
            'last_date': None
        }

        try:
            self.logger.debug(f"更新股票 {symbol} (重试 {retry_count + 1}/{self.update_config['max_retries']})")

            # 获取现有数据的最后日期
            info = self.data_storage.get_data_info(symbol, self.update_config['interval'])

            if info['file_exists'] and info['record_count'] > 0:
                # 增量更新：现有数据存在
                existing_data = self.data_storage.load_data(symbol, self.update_config['interval'])

                if existing_data.empty:
                    result['error'] = "现有数据为空"
                    return result

                # 确保Date列存在（可能是索引或列）
                if 'Date' not in existing_data.columns:
                    if existing_data.index.name == 'Date':
                        # 如果Date是索引，转换为列
                        existing_data = existing_data.reset_index()
                    else:
                        result['error'] = "数据格式错误：缺少Date列"
                        return result

                # 确定更新开始日期
                last_date = pd.to_datetime(existing_data['Date'].max())
                start_date = (last_date + timedelta(days=1)).strftime('%Y-%m-%d')
                end_date = datetime.now().strftime('%Y-%m-%d')

                # 如果已经是最新的，跳过
                if start_date > end_date:
                    result['success'] = True
                    result['last_date'] = last_date.strftime('%Y-%m-%d')
                    return result

                # 获取新数据（使用数据源管理器的回退机制）
                new_data, source_name = self.data_source_manager.fetch_data_with_fallback(
                    symbol,
                    self.update_config['interval'],
                    start_date,
                    end_date
                )
                if source_name:
                    self.logger.debug(f"使用数据源 {source_name} 获取 {symbol} 的新数据")

                if new_data.empty:
                    result['success'] = True
                    result['last_date'] = last_date.strftime('%Y-%m-%d')
                    return result

                # 确保新数据也有Date列（可能是索引或列）
                if 'Date' not in new_data.columns:
                    if new_data.index.name == 'Date':
                        # 如果Date是索引，转换为列
                        new_data = new_data.reset_index()
                    else:
                        result['error'] = "新数据格式错误：缺少Date列"
                        return result

                # 合并数据
                combined_data = pd.concat([existing_data, new_data], ignore_index=True)
                combined_data = combined_data.drop_duplicates(subset=['Date'], keep='last')
                combined_data = combined_data.sort_values('Date').reset_index(drop=True)

                # 保存数据
                success = self.data_storage.save_data(symbol, combined_data, self.update_config['interval'])

                if success:
                    result['success'] = True
                    result['new_records'] = len(new_data)
                    result['last_date'] = pd.to_datetime(combined_data['Date'].max()).strftime('%Y-%m-%d')
                    self.logger.info(f"成功更新 {symbol}: {len(new_data)} 条新记录")
                else:
                    result['error'] = "保存失败"

            else:
                # 全量更新：文件不存在或为空，获取全部历史数据
                self.logger.info(f"股票 {symbol} 没有历史数据，开始获取全部历史数据")

                # 确定合适的开始日期（根据股票类型）
                if symbol.startswith('688'):  # 科创板
                    start_date = '2019-07-22'  # 科创板开市日期
                elif symbol.startswith('300'):  # 创业板
                    start_date = '2009-10-30'  # 创业板开市日期
                elif symbol.startswith('688'):  # 科创板
                    start_date = '2019-07-22'
                elif symbol.startswith('000') or symbol.startswith('001') or symbol.startswith('002'):  # 深圳主板
                    start_date = '1990-12-19'
                elif symbol.startswith('600') or symbol.startswith('601') or symbol.startswith('603'):  # 上海主板
                    start_date = '1990-12-19'
                else:
                    start_date = '2000-01-01'  # 默认开始日期

                end_date = datetime.now().strftime('%Y-%m-%d')

                # 获取全部历史数据（使用数据源管理器的回退机制）
                historical_data, source_name = self.data_source_manager.fetch_data_with_fallback(
                    symbol,
                    self.update_config['interval'],
                    start_date,
                    end_date
                )
                if source_name:
                    self.logger.debug(f"使用数据源 {source_name} 获取 {symbol} 的历史数据")

                if historical_data.empty:
                    result['error'] = "未能获取历史数据"
                    return result

                # 确保Date列存在（可能是索引或列）
                if 'Date' not in historical_data.columns:
                    if historical_data.index.name == 'Date':
                        # 如果Date是索引，转换为列
                        historical_data = historical_data.reset_index()
                    else:
                        result['error'] = "数据格式错误：缺少Date列"
                        return result

                # 保存数据
                success = self.data_storage.save_data(symbol, historical_data, self.update_config['interval'])

                if success:
                    result['success'] = True
                    result['new_records'] = len(historical_data)
                    result['last_date'] = pd.to_datetime(historical_data['Date'].max()).strftime('%Y-%m-%d')
                    self.logger.info(f"成功获取 {symbol} 历史数据: {len(historical_data)} 条记录")
                else:
                    result['error'] = "保存失败"

        except Exception as e:
            result['error'] = str(e)
            self.logger.error(f"更新 {symbol} 失败 (重试 {retry_count + 1}): {e}")

        return result

    def load_failed_stocks(self) -> List[str]:
        """加载之前失败的股票列表"""
        try:
            failed_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'failed_stocks.json'
            if not failed_file.exists():
                return []

            with open(failed_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                return data.get('failed_symbols', [])

        except Exception as e:
            self.logger.error(f"加载失败股票列表失败: {e}")
            return []

    def save_failed_stocks(self, failed_symbols: List[str]):
        """保存失败的股票列表"""
        try:
            failed_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'failed_stocks.json'
            data = {
                'timestamp': datetime.now().isoformat(),
                'failed_symbols': failed_symbols,
                'count': len(failed_symbols)
            }

            with open(failed_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)

        except Exception as e:
            self.logger.error(f"保存失败股票列表失败: {e}")

    def save_update_report(self, results: List[Dict[str, Any]]):
        """保存更新报告"""
        try:
            report_data = {
                'timestamp': datetime.now().isoformat(),
                'config': self.update_config,
                'summary': {
                    'total': len(results),
                    'success': len([r for r in results if r['success']]),
                    'failed': len([r for r in results if not r['success']]),
                    'total_new_records': sum(r.get('new_records', 0) for r in results)
                },
                'results': results,
                'failed_symbols': [r['symbol'] for r in results if not r['success']]
            }

            # 保存详细报告
            report_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'update_report.json'
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump(report_data, f, ensure_ascii=False, indent=2)

            # 保存简要报告
            summary_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'update_summary.txt'
            with open(summary_file, 'w', encoding='utf-8') as f:
                f.write(f"更新时间: {report_data['timestamp']}\n")
                f.write(f"更新配置: {self.update_config['interval']}, 最近{self.update_config['days_to_update']}天\n")
                f.write(f"总股票数: {report_data['summary']['total']}\n")
                f.write(f"成功更新: {report_data['summary']['success']}\n")
                f.write(f"更新失败: {report_data['summary']['failed']}\n")
                f.write(f"新增记录: {report_data['summary']['total_new_records']}\n")
                f.write(f"成功率: {report_data['summary']['success']/report_data['summary']['total']*100:.1f}%\n")

                if report_data['failed_symbols']:
                    f.write(f"\n失败的股票: {report_data['failed_symbols'][:10]}")

            self.logger.info(f"更新报告已保存到: {report_file}")

        except Exception as e:
            self.logger.error(f"保存更新报告失败: {e}")

    def run_daily_update(self, symbols: List[str] = None, update_mode: str = "incremental"):
        """运行每日更新"""
        try:
            start_time = datetime.now()
            self.logger.info("=" * 60)
            mode_text = "全量更新" if update_mode == "full" else "增量更新"
            self.logger.info(f"开始A股{mode_text}")
            self.logger.info("=" * 60)

            # 获取股票列表
            if symbols is None:
                symbols = self.get_update_stock_list(update_mode)

            if not symbols:
                self.logger.info("没有需要更新的股票")
                return True

            # 加载之前失败的股票
            if self.update_config['update_failed_stocks']:
                failed_symbols = self.load_failed_stocks()
                if failed_symbols:
                    self.logger.info(f"加载了 {len(failed_symbols)} 只之前失败的股票")
                    symbols = list(set(symbols + failed_symbols))

            total_count = len(symbols)
            self.logger.info(f"总共需要更新 {total_count} 只股票")

            # 分批处理
            batch_size = self.update_config['batch_size']
            total_batches = (total_count + batch_size - 1) // batch_size

            all_results = []

            for batch_idx in range(total_batches):
                batch_start = batch_idx * batch_size
                batch_end = min(batch_start + batch_size, total_count)
                batch_symbols = symbols[batch_start:batch_end]

                self.logger.info(f"处理批次 {batch_idx + 1}/{total_batches}: {len(batch_symbols)} 只股票")

                batch_results = []
                for i, symbol in enumerate(batch_symbols):
                    self.logger.debug(f"更新 {symbol} ({batch_start + i + 1}/{total_count})")

                    # 更新股票数据
                    result = self.update_stock(symbol)
                    batch_results.append(result)

                    # 股票间延迟
                    time.sleep(self.update_config['delay_between_stocks'])

                all_results.extend(batch_results)

                # 批次间延迟
                if batch_idx < total_batches - 1:
                    time.sleep(self.update_config['delay_between_batches'])

                # 批次统计
                batch_success = len([r for r in batch_results if r['success']])
                batch_failed = len(batch_results) - batch_success
                self.logger.info(f"批次 {batch_idx + 1} 完成: 成功 {batch_success}, 失败 {batch_failed}")

            # 保存失败的股票列表
            failed_symbols = [r['symbol'] for r in all_results if not r['success']]
            if failed_symbols:
                self.save_failed_stocks(failed_symbols)

            # 保存更新报告
            if self.update_config['save_report']:
                self.save_update_report(all_results)

            # 最终统计
            end_time = datetime.now()
            duration = end_time - start_time

            success_count = len([r for r in all_results if r['success']])
            failed_count = len(all_results) - success_count
            total_new_records = sum(r.get('new_records', 0) for r in all_results)

            self.logger.info("=" * 60)
            self.logger.info("更新完成！")
            self.logger.info("=" * 60)
            self.logger.info(f"总耗时: {duration}")
            self.logger.info(f"总股票数: {total_count}")
            self.logger.info(f"成功更新: {success_count}")
            self.logger.info(f"更新失败: {failed_count}")
            self.logger.info(f"新增记录: {total_new_records}")
            self.logger.info(f"成功率: {success_count / total_count * 100:.1f}%")

            if failed_symbols:
                self.logger.warning(f"失败的股票: {failed_symbols[:10]}...")

            return failed_count == 0

        except Exception as e:
            self.logger.error(f"更新过程出现异常: {e}")
            return False

    def run_quick_update(self, sample_count: int = 50):
        """快速更新（用于测试）"""
        try:
            symbols = self.get_update_stock_list()
            sample_symbols = symbols[:sample_count]

            self.logger.info(f"运行快速更新，处理 {len(sample_symbols)} 只股票")

            return self.run_daily_update(sample_symbols)

        except Exception as e:
            self.logger.error(f"快速更新失败: {e}")
            return False

    def check_update_status(self) -> Dict[str, Any]:
        """检查更新状态"""
        try:
            symbols = self.get_update_stock_list()
            total_count = len(symbols)

            outdated_count = 0
            up_to_date_count = 0
            failed_count = 0
            outdated_stocks = []

            for symbol in symbols:
                try:
                    info = self.data_storage.get_data_info(symbol, '1d')
                    if info['file_exists']:
                        existing_data = self.data_storage.load_data(symbol, '1d')
                        if not existing_data.empty:
                            # 确保Date列存在
                            if 'Date' not in existing_data.columns:
                                failed_count += 1
                                continue

                            last_date = pd.to_datetime(existing_data['Date'].max())
                            days_diff = (datetime.now() - last_date).days

                            if days_diff > 2:  # 超过2天未更新
                                outdated_count += 1
                                outdated_stocks.append(symbol)
                            else:
                                up_to_date_count += 1
                        else:
                            failed_count += 1
                    else:
                        failed_count += 1
                except Exception as e:
                    self.logger.warning(f"检查 {symbol} 状态失败: {e}")
                    failed_count += 1

            return {
                'total_stocks': total_count,
                'up_to_date': up_to_date_count,
                'outdated': outdated_count,
                'failed': failed_count,
                'update_rate': up_to_date_count / total_count * 100 if total_count > 0 else 0,
                'sample_outdated': outdated_stocks[:10]  # 显示前10个过期的
            }

        except Exception as e:
            self.logger.error(f"检查更新状态失败: {e}")
            return {}


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='A股每日数据更新')
    parser.add_argument('--mode', choices=['full', 'quick', 'status'], default='full',
                       help='运行模式: full=全量更新, quick=快速更新, status=检查状态')
    parser.add_argument('--sample', type=int, default=50,
                       help='快速更新的样本数量')
    parser.add_argument('--symbols', nargs='+', help='指定要更新的股票代码')

    args = parser.parse_args()

    updater = AShareUpdater()

    try:
        if args.mode == 'status':
            status = updater.check_update_status()
            print("\n=== A股更新状态 ===")
            print(f"总股票数: {status['total_stocks']}")
            print(f"已最新: {status['up_to_date']}")
            print(f"需要更新: {status['outdated']}")
            print(f"异常数量: {status['failed']}")
            print(f"更新率: {status['update_rate']:.1f}%")

            if status['sample_outdated']:
                print(f"\n需要更新的股票示例: {status['sample_outdated']}")

        elif args.mode == 'quick':
            print(f"运行快速更新，样本数量: {args.sample}")
            success = updater.run_quick_update(args.sample)
            sys.exit(0 if success else 1)

        elif args.mode == 'full':
            if args.symbols:
                print(f"更新指定股票: {args.symbols}")
                success = updater.run_daily_update(args.symbols)
            else:
                print("运行全量更新")
                success = updater.run_daily_update()
            sys.exit(0 if success else 1)

    except KeyboardInterrupt:
        print("\n用户中断更新过程")
        sys.exit(1)
    except Exception as e:
        print(f"更新过程出现错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()